import json
import logging
from typing import List, Dict, Any, Set, Tuple, Optional

# 尝试导入依赖，如果环境没有这些库，RAG功能将无法使用
try:
    import jieba
    from rank_bm25 import BM25Okapi
    HAS_RAG_DEPS = True
except ImportError:
    HAS_RAG_DEPS = False

logger = logging.getLogger(__name__)

class RagHandler:
    """
    RAG (检索增强生成) 处理类
    
    用于处理超长文档，通过检索相关切片（页）来减少输入上下文
    """
    
    def __init__(self, max_length: int = 100000, window_size: int = 1):
        """
        初始化 RAG 处理器
        
        Args:
            max_length: 触发 RAG 的内容长度阈值（字符数）。
                        DeepSeek V3 支持 128k token，约等于 10-20w 中文字符。
                        设置为 100000 比较安全。
            window_size: 命中页面的上下文窗口大小（前后各取多少页）
        """
        self.max_length = max_length
        self.window_size = window_size
        
        if not HAS_RAG_DEPS:
            logger.warning("缺少 RAG 依赖 (jieba, rank_bm25)，将无法使用检索功能。请运行 `pip install jieba rank_bm25` 安装。")
            
    def should_use_rag(self, content_length: int) -> bool:
        """判断是否需要使用 RAG"""
        if not HAS_RAG_DEPS:
            return False
        return content_length > self.max_length
        
    def get_relevant_content(
        self, 
        parse_results: List[Any], 
        query: str,
        top_k: int = 5
    ) -> Tuple[str, str]:
        """
        获取相关的文档内容（Markdown 和 Middle JSON）
        
        Args:
            parse_results: 解析结果列表 (ParseResult 对象列表)
            query: 查询语句 (通常是 字段名 + 提取规则)
            top_k: 检索最相关的页数
            
        Returns:
            Tuple[str, str]: (相关的 Markdown 内容, 相关的 Middle JSON 字符串)
        """
        if not parse_results:
            return "", "[]"
            
        # 1. 提取所有页面文本，建立索引
        all_pages = []
        corpus = []
        
        global_page_idx = 0
        
        for file_idx, result in enumerate(parse_results):
            if not hasattr(result, 'pages'):
                continue
                
            for page in result.pages:
                # 提取页面文本（用于检索）
                page_text = self._extract_page_text(page)
                
                # 存储页面信息
                all_pages.append({
                    "file_idx": file_idx,
                    "page_obj": page,
                    "global_page_idx": global_page_idx,
                    "text": page_text
                })
                
                # 分词用于 BM25
                tokens = list(jieba.cut(page_text))
                corpus.append(tokens)
                
                global_page_idx += 1
                
        if not corpus:
            return "", "[]"
            
        # 2. BM25 检索
        bm25 = BM25Okapi(corpus)
        tokenized_query = list(jieba.cut(query))
        
        # 获取分数并选出 Top K
        doc_scores = bm25.get_scores(tokenized_query)
        top_n_indices = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True)[:top_k]
        
        # 3. 扩展上下文窗口 (前后 window_size 页)
        selected_indices = set()
        for idx in top_n_indices:
            # 只有当分数大于0时才认为有效（避免无关内容）
            if doc_scores[idx] > 0:
                start = max(0, idx - self.window_size)
                end = min(len(all_pages), idx + self.window_size + 1)
                for i in range(start, end):
                    selected_indices.add(i)
        
        if not selected_indices:
            # 如果没有匹配到任何内容（分数全为0），回退到只取前几页或报错？
            # 这里选择取前 5 页作为兜底
            logger.warning(f"RAG 检索未匹配到相关内容: {query}，使用前 5 页兜底")
            selected_indices = set(range(min(5, len(all_pages))))
            
        # 4. 重新组织选中的页面
        sorted_indices = sorted(list(selected_indices))
        selected_pages = [all_pages[i] for i in sorted_indices]
        
        logger.info(f"RAG 检索命中 {len(selected_indices)}/{len(all_pages)} 页 (Top {top_k} + Window {self.window_size})")
        
        # 5. 生成 Markdown 和 Middle JSON
        return self._generate_outputs(selected_pages, parse_results)

    def _extract_page_text(self, page: Any) -> str:
        """从页面对象提取纯文本"""
        texts = []
        # 尝试从 contents 提取
        if hasattr(page, 'contents'):
            for content in page.contents:
                if hasattr(content, 'text') and content.text:
                    texts.append(content.text)
        # 兜底：如果 contents 为空，尝试 layout_blocks
        elif hasattr(page, 'layout_blocks'):
            for block in page.layout_blocks:
                if hasattr(block, 'text') and block.text:
                    texts.append(block.text)
        
        return "\n".join(texts)

    def _generate_outputs(self, selected_pages: List[Dict], parse_results: List[Any]) -> Tuple[str, str]:
        """生成 Markdown 和 Middle JSON"""
        md_parts = []
        middle_data = []
        
        current_file_idx = -1
        
        for item in selected_pages:
            file_idx = item['file_idx']
            page = item['page_obj']
            global_page_idx = item['global_page_idx']
            
            # Markdown 处理：添加文件分隔符（如果是新文件）
            if file_idx != current_file_idx:
                file_name = getattr(parse_results[file_idx], 'file_name', f'文件{file_idx+1}')
                md_parts.append(f"\n# === {file_name} (Page {page.page_num}) ===\n")
                current_file_idx = file_idx
            else:
                md_parts.append(f"\n## Page {page.page_num}\n")
            
            # 重建该页的 Markdown
            page_md = self._reconstruct_page_markdown(page)
            md_parts.append(page_md)
            
            # Middle JSON 处理
            # 复用 ExtractionService 的逻辑结构，但只包含选中的页
            # 需要构造符合 middle.json 结构的 dict
            # 注意：middle.json 依赖 page_idx，这里我们保持 global_page_idx 以便溯源能对应上
            page_data = self._construct_middle_page(page, global_page_idx)
            middle_data.append(page_data)
            
        return "\n".join(md_parts), json.dumps(middle_data, ensure_ascii=False)

    def _reconstruct_page_markdown(self, page: Any) -> str:
        """重建单页 Markdown"""
        # 这里做一个简单的重建，按顺序拼接 contents
        lines = []
        if hasattr(page, 'contents'):
            for content in page.contents:
                text = getattr(content, 'text', '')
                type_ = getattr(content, 'type', 'text')
                
                if not text:
                    continue
                    
                if type_ in ['title', 'header']:
                    lines.append(f"## {text}")
                elif type_ == 'list':
                    lines.append(f"- {text}")
                elif type_ == 'table':
                    # 如果有 table_body (HTML)，优先使用
                    table_body = getattr(content, 'table_body', '')
                    if table_body:
                        lines.append(table_body)
                    else:
                        lines.append(text)
                else:
                    lines.append(text)
                    
        return "\n\n".join(lines)

    def _construct_middle_page(self, page: Any, global_page_idx: int) -> Dict:
        """构造单页 middle json 数据"""
        # 这部分逻辑模仿 ExtractionService._extract_middle_from_single 中的逻辑
        # 但简化为直接从 page 对象提取
        
        layout_dets = []
        if hasattr(page, 'layout_blocks'):
            for block in page.layout_blocks:
                # 处理 bbox
                bbox = [0, 0, 0, 0]
                if hasattr(block, 'bbox'):
                    b = block.bbox
                    if hasattr(b, 'x0'):
                         bbox = [b.x0, b.y0, b.x1, b.y1]
                    elif isinstance(b, (list, tuple)) and len(b) == 4:
                        bbox = list(b)
                
                # 构建 lines (为了溯源高亮)
                lines_data = []
                if hasattr(block, 'lines'):
                    for line in block.lines:
                        line_bbox = [0,0,0,0] # 简化，如果有真实数据则使用
                        if hasattr(line, 'bbox'):
                             lb = line.bbox
                             if isinstance(lb, (list, tuple)) and len(lb) == 4: line_bbox = list(lb)
                        
                        spans_data = []
                        if hasattr(line, 'spans'):
                            for span in line.spans:
                                span_text = getattr(span, 'text', '')
                                if span_text:
                                    span_bbox = [0,0,0,0]
                                    if hasattr(span, 'bbox'):
                                        sb = span.bbox
                                        if isinstance(sb, (list, tuple)) and len(sb) == 4: span_bbox = list(sb)
                                    
                                    spans_data.append({
                                        "bbox": span_bbox,
                                        "content": span_text,
                                        "type": getattr(span, 'type', 'text')
                                    })
                        
                        if spans_data:
                            lines_data.append({
                                "bbox": line_bbox,
                                "spans": spans_data
                            })
                
                layout_dets.append({
                    "bbox": bbox,
                    "text": getattr(block, 'text', ''),
                    "category_type": getattr(block, 'category_type', 'text'),
                    "lines": lines_data
                })
                
        return {
            "page_idx": global_page_idx,
            "layout_dets": layout_dets
        }


